package com.dbl.ml.integration.lucene;

import com.dbl.ml.string.StringUtils;
import com.google.common.collect.Maps;
import org.apache.lucene.analysis.*;
import org.apache.lucene.analysis.standard.StandardAnalyzer;
import org.apache.lucene.index.Term;
import org.apache.lucene.queryParser.ParseException;
import org.apache.lucene.queryParser.QueryParser;
import org.apache.lucene.search.*;
import org.apache.lucene.search.spans.SpanNearQuery;
import org.apache.lucene.search.spans.SpanQuery;
import org.apache.lucene.search.spans.SpanTermQuery;

import java.io.IOException;
import java.io.StringReader;
import java.util.Map;

/**
 * Author: Yura Korolov
 * Date: 14.12.2010
 * Time: 12:54:22
 */
public abstract class LuceneQueryBuilder {
    public static RootQueryBuilder create() {
        return new RootQueryBuilderImpl();
    }

    public static enum AnalyzerType {
        STD(new StandardAnalyzer()),
        WHITeSPACE(new WhitespaceAnalyzer());

        Analyzer an;

        AnalyzerType(Analyzer an) {
            this.an = an;
        }
    }


    public interface QueryBuilder<T extends Query> {
        T toQuery();

        QueryBuilder<T> boost(float value);

    }

    public interface RootQueryBuilder {
        RootQueryBuilder setAnalyzer(Analyzer an);

        RootQueryBuilder setAnalyzer(String field, Analyzer an);

        RootQueryBuilder setAnalyzer(String field, AnalyzerType an);

        TermQuery term(String field, String term);
        RootQueryBuilder setAnalyzeTerms(boolean flag);

        BooleanQueryBuilder bool();

        Query phrase(Term... terms);

        ParseQueryBuilder parseQuery(String field, String query);


        SpanNearQuery spanNear(int slop, boolean inOrder, SpanQuery... clauses);

        SpanTermQuery spanTerm(String field, String term);

        Query boost(Query q, float boost);


    }

    public static class BooleanQueryBuilder implements QueryBuilder<Query> {
        BooleanQuery bq = new BooleanQuery();

        public BooleanQueryBuilder should(Query query) {
            bq.add(query, BooleanClause.Occur.SHOULD);
            return this;
        }

        public BooleanQueryBuilder must(Query query) {
            bq.add(query, BooleanClause.Occur.MUST);
            return this;
        }

        public BooleanQueryBuilder mustNot(Query query) {
            bq.add(query, BooleanClause.Occur.MUST_NOT);
            return this;
        }

        public BooleanQueryBuilder setMinimumNumberShouldMatch(int min) {
            bq.setMinimumNumberShouldMatch(min);
            return this;
        }

        public Query toQuery() {
            return bq;
        }

        public QueryBuilder<Query> boost(float value) {
            bq.setBoost(value);
            return this;
        }
    }

    public static class ParseQueryBuilder implements QueryBuilder<Query> {
        String field;
        String query;
        RootQueryBuilderImpl root;
        boolean escaped;
        boolean defaultAnd = true;
        float boost = 1f;
        boolean allOnBlank = false;

        public ParseQueryBuilder(String field, RootQueryBuilderImpl root, String query) {
            this.field = field;
            this.root = root;
            this.query = query;
        }

        public ParseQueryBuilder escape() {
            this.escaped = true;
            return this;
        }

        public ParseQueryBuilder defaultAnd() {
            defaultAnd = true;
            return this;
        }

        public ParseQueryBuilder defaultOr() {
            defaultAnd = false;
            return this;
        }

        public ParseQueryBuilder allOnBlank() {
            this.allOnBlank = true;
            return this;
        }

        public Query toQuery() {
            try {
                if (allOnBlank && StringUtils.isBlank(query)) return new MatchAllDocsQuery();
                if (!allOnBlank && StringUtils.isBlank(query)) return new MatchAllDocsQuery();
                QueryParser qp = new QueryParser(field, root.analyzer());
                qp.setDefaultOperator(defaultAnd ? QueryParser.Operator.AND : QueryParser.Operator.OR);
                if(escaped) {
                    query = QueryParser.escape(query);
                    query =  query.replace("AND", "and").replace("OR", "or");
                }
                Query query1 = qp.parse(query);
                query1.setBoost(boost);
                return query1;
            } catch (ParseException e) {
                throw new RuntimeException(e);
            }
        }

        public QueryBuilder<Query> boost(float value) {
            boost = value;
            return this;
        }
    }

    public static class RootQueryBuilderImpl implements RootQueryBuilder {

        PerFieldAnalyzerWrapper perFieldAnalyzer = null;
        Map<String, Analyzer> anMap = Maps.newHashMap();
        Analyzer defaultAnalyzer = AnalyzerType.STD.an;
        private boolean tokenizeTerms = false;

        public RootQueryBuilder setAnalyzer(Analyzer an) {
            defaultAnalyzer = an;
            this.perFieldAnalyzer = null;
            return this;
        }

        public RootQueryBuilder setAnalyzer(String field, Analyzer an) {
            anMap.put(field, an);
            if (defaultAnalyzer == null) defaultAnalyzer = an;
            this.perFieldAnalyzer = null;
            return this;
        }

        public RootQueryBuilder setAnalyzer(String field, AnalyzerType an) {
            return setAnalyzer(field, an.an);
        }

        public TermQuery term(String field, String term) {
            return new TermQuery(newTerm(field, term));
        }

        public RootQueryBuilder setAnalyzeTerms(boolean flag) {
            tokenizeTerms = flag;
            return this;
        }

        private Term newTerm(String field, String term) {
            if (tokenizeTerms) {
                TokenStream stream = analyzer().tokenStream(field, new StringReader(term));
                try {
                    Token token = stream.next();
                    term = token.term();
                } catch (IOException e) {
                    throw new RuntimeException(e);
                }
            }
            return new Term(field, term);
        }

        public BooleanQueryBuilder bool() {
            return new BooleanQueryBuilder();
        }

        public PhraseQuery phrase(Term... terms) {
            PhraseQuery pq = new PhraseQuery();
            for (Term t : terms) pq.add(t);
            return pq;
        }

        public ParseQueryBuilder parseQuery(String field, String query) {
            return new ParseQueryBuilder(field, this, query);
        }

        public SpanNearQuery spanNear(int slop, boolean inOrder, SpanQuery... clauses) {
            return new SpanNearQuery(clauses, slop, inOrder);
        }

        public SpanTermQuery spanTerm(String field, String term) {
            return new SpanTermQuery(newTerm(field, term));
        }

        public Query boost(Query q, float boost) {
            q.setBoost(boost);
            return q;
        }


        Analyzer analyzer() {
            if (perFieldAnalyzer != null) {
                return perFieldAnalyzer;
            }
            perFieldAnalyzer = new PerFieldAnalyzerWrapper(defaultAnalyzer);
            for (Map.Entry<String, Analyzer> anE : anMap.entrySet()) {
                perFieldAnalyzer.addAnalyzer(anE.getKey(), anE.getValue());
            }
            return perFieldAnalyzer;
        }
    }

    public static void main(String... args) {
        /*
        qb.bool()
            .should()
            .must()
            .mustNot()
            .must( qb.spanNear(qb.near()  ))

            .toQuery()


        qb.bool()
         */

    }

}


